# encoding=utf-8
import threading
import cv2
import numpy
import onnxruntime as ort
from typing import List
from transform import center_crop, resized_edge


class ConvNeXtWrapper:
    _inst = None
    _lock = threading.Lock()

    def __init__(self, model_path: str):
        """运行眼底疾病诊断，可实现青光眼二分类和糖尿病视网膜病变5分级.
              * glaucoma (int): 1代表可疑青光眼，0代表非青光眼
              * dr_level (int): 糖尿病视网膜病变的等级，0代表健康，4代表患病程度最严重
                                0 为非糖尿病视网膜病变
                                1 代表轻度非增生性的糖尿病视网膜病变
                                2 中度
                                3 重度
                                4 增生性糖尿病视网膜病变
              * amd (int): 1代表可疑年龄相关性黄斑变性，0代表非黄斑变性
              * pm (int): 1代表病理性近视，0代表非病理性近视
        """
        self.model = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
        self.input_name = [ele.name for _, ele in enumerate(self.model.get_inputs())]
        self.output_name = [ele.name for _, ele in enumerate(self.model.get_outputs())]

    def _preprocess(self, data_in: str):
        """
            bug fix:
                由于官方子项目源码和顶层项目代码不一致，
                结果也不一致, 所以将代码修改到顶层项目代码
        """
        img = cv2.imread(data_in)
        img = resized_edge(img, 448, edge='long')
        img = center_crop(img, 448)
        mean = [0.48145466 * 255, 0.4578275 * 255, 0.40821073 * 255]
        std = [0.26862954 * 255, 0.26130258 * 255, 0.27577711 * 255]
        img = (img - mean) / std
        img = img[..., ::-1]  # bgr to rgb
        img = img.transpose((2, 0, 1))
        img = img.astype('float32')
        img = img[numpy.newaxis, ...]
        return img

    def _postprocess(self, data_list: List[numpy.array]):
        glaucoma = data_list[0][0].argmax()
        dr = data_list[1][0].argmax()
        amd = data_list[2][0].argmax()
        pm = data_list[3][0].argmax()
        return glaucoma, dr, amd, pm

    def infer(self, image_path: str):
        trans_data = self._preprocess(image_path)
        res_list = self.model.run(self.output_name, {self.input_name[0]: trans_data})
        return self._postprocess(res_list)

    def __new__(cls, *args, **kwargs):
        ConvNeXtWrapper._lock.acquire()
        if cls._inst is None:
            cls._inst = super(ConvNeXtWrapper, cls).__new__(cls)
        ConvNeXtWrapper._lock.release()
        return cls._inst


if __name__ == '__main__':
    cov = ConvNeXtWrapper("convnext-tiny.onnx")
    print(cov.infer("1.jpg"))

